package edu.ynu.se.xiecheng.achitectureclass.auth.config;

import edu.ynu.se.xiecheng.achitectureclass.auth.dao.UserDAO;
import edu.ynu.se.xiecheng.achitectureclass.auth.entity.User;
import edu.ynu.se.xiecheng.achitectureclass.auth.token.BearerToken;
import org.apache.shiro.authc.*;
import org.apache.shiro.authz.AuthorizationInfo;
import org.apache.shiro.authz.SimpleAuthorizationInfo;
import org.apache.shiro.realm.AuthorizingRealm;
import org.apache.shiro.session.Session;
import org.apache.shiro.session.mgt.DefaultSessionKey;
import org.apache.shiro.SecurityUtils;
import org.apache.shiro.subject.PrincipalCollection;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.security.crypto.password.PasswordEncoder;

import java.util.stream.Collectors;

public class CustomRealm extends AuthorizingRealm {

    @Autowired
    private UserDAO userDAO;

    @Autowired
    private PasswordEncoder passwordEncoder;

    @Override
    public boolean supports(AuthenticationToken token) {
        return token instanceof UsernamePasswordToken || token instanceof BearerToken;
    }

    @Override
    protected AuthorizationInfo doGetAuthorizationInfo(PrincipalCollection principals) {
        String username = (String) principals.getPrimaryPrincipal();
        SimpleAuthorizationInfo authorizationInfo = new SimpleAuthorizationInfo();
        
        userDAO.findByUsername(username).ifPresent(user -> {
            authorizationInfo.setRoles(user.getRoles().stream()
                    .map(role -> role.getName())
                    .collect(Collectors.toSet()));
        });
        
        return authorizationInfo;
    }

    @Override
    protected AuthenticationInfo doGetAuthenticationInfo(AuthenticationToken token) throws AuthenticationException {
        if (token instanceof BearerToken) {
            // 处理 Bearer Token
            String sessionId = (String) token.getPrincipal();
            try {
                Session session = SecurityUtils.getSecurityManager().getSession(new DefaultSessionKey(sessionId));
                if (session != null) {
                    String username = (String) session.getAttribute("username");
                    if (username != null) {
                        // 验证用户是否存在
                        User user = userDAO.findByUsername(username)
                                .orElseThrow(() -> new UnknownAccountException("用户不存在"));
                        return new SimpleAuthenticationInfo(username, sessionId, getName());
                    }
                }
                throw new AuthenticationException("Invalid session");
            } catch (Exception e) {
                throw new AuthenticationException("Invalid token: " + e.getMessage());
            }
        } else {
            // 处理用户名密码认证
            String username = (String) token.getPrincipal();
            String password = new String((char[]) token.getCredentials());

            // 从数据库查找用户
            User user = userDAO.findByUsername(username)
                    .orElseThrow(() -> new UnknownAccountException("用户不存在"));

            // 验证密码
            if (!passwordEncoder.matches(password, user.getPassword())) {
                throw new IncorrectCredentialsException("密码错误");
            }

            // 返回认证信息
            return new SimpleAuthenticationInfo(username, password, getName());
        }
    }
} 